package promote.brozen.algorithm.tree.binary.avl;

import promote.brozen.algorithm.tree.Tree;
import promote.brozen.algorithm.tree.binary.RecursionBinaryTreeWalker;

import java.util.function.Consumer;

/**
 * 平衡二叉树，又名为AVL树
 * 递归的思想
 *
 * 树具有自相似性，所以用递归可以解决？
 *
 * @author Brozen
 * @date 2019/12/24 10:21 AM
 */
public class AVLTree<K extends Comparable<K>> implements Tree<K> {

    private AVLTreeNode<K> root;

    @Override
    public void add(K key) {
        this.root = doAdd(this.root, key);
    }

    /**
     * 在以入参节点为根节点的子树中，插入指定的key，并返回经过平衡后的新的子树根节点
     */
    private AVLTreeNode<K> doAdd(AVLTreeNode<K> node, K key) {
        if (node == null) {
            return new AVLTreeNode<>(key);
        }

        // 先执行插入
        int compareResult;
        if ((compareResult = node.getKey().compareTo(key)) > 0) {
            node.setLeftChild(doAdd(node.getLeftChild(), key));
        } else if (compareResult < 0) {
            node.setRightChild(doAdd(node.getRightChild(), key));
        } else {
            // 如果没有插入新节点，是更新操作，不需要平衡
            node.setKey(key);
            return node;
        }

        // 如果插入了新的节点，则需要进行平衡
        return doBalance(node);

    }

    @Override
    public void remove(K key) {
        this.root = doRemove(this.root, key);
    }

    /**
     * 在以入参节点为根节点的子树中，删除指定的key，并返回经过平衡后的新的子树根节点
     */
    private AVLTreeNode<K> doRemove(AVLTreeNode<K> node, K key) {
        if (node == null) {
            return null;
        }

        int compareResult;
        if ((compareResult = node.getKey().compareTo(key)) > 0) {
            node.setLeftChild(doRemove(node.getLeftChild(), key));
        } else if (compareResult < 0) {
            node.setRightChild(doRemove(node.getRightChild(), key));
        } else {
            // 要删除这个节点了，根据二叉查找树的规则，删除节点后根据节点的子节点情况进行处理
            if (node.getLeftChild() == null) {
                // 左子节点为空，则用右子节点代替当前节点
                node = node.getRightChild();
            } else if (node.getRightChild() == null) {
                // 右子节点为空，则用左子节点代替当前节点
                node = node.getLeftChild();
            } else {
                // 左右子节点均存在，则取右子树的最小节点为后继节点代替当前节点
                AVLTreeNode<K> successor = getMin(node.getRightChild());
                // 后继节点不是当前节点的右子节点时，需要用后继节点的右子节点代替后继节点，
                //   然后将后继节点的右子节点指向当前节点的右子节点，最后将后继节点的左子节点
                //   指向当前节点的左子节点
                // 后继节点是当前节点的右子节点时，则后继节点的子节点关系不需要变化，只要
                //   把后继节点的左子节点指向当前节点的左子节点就行
                if (successor != node.getRightChild()) {
                    successor.getParent().setLeftChild(successor.getRightChild());
                    successor.setRightChild(node.getRightChild());
                }
                successor.setLeftChild(node.getLeftChild());

                // 此时当前节点node已经被从树中删除了，接下来将后继节点作为当前节点，进行平衡操作
                node = successor;
            }
        }

        // 删除后，当前节点node有可能为空，如果当前节点为空，
        // 说明删除的是叶子节点，不需要再平衡当前节点
        if (node == null) {
            return null;
        }

        // 当前节点不为空，则说明以当前节点为根的子树中节点发生了变化，需要进行平衡
        return doBalance(node);
    }

    /**
     * 获取以入参节点为根节点的子树中key值最小的节点
     */
    private AVLTreeNode<K> getMin(AVLTreeNode<K> node) {
        if (node == null) {
            return null;
        }

        while (node.getLeftChild() != null) {
            node = node.getLeftChild();
        }

        return node;
    }

    /**
     * 平衡以入参节点为根节点的树，并返回新的根节点
     */
    private AVLTreeNode<K> doBalance(AVLTreeNode<K> node) {
        // 平衡之前重新计算节点高度
        // 需要注意的是，旋转是树下层到根执行的，所以这里计算高度也会算从下往上计算，
        // 而节点高度的定义是从叶子节点到节点路径的最大值，则下层高度是对的的话，上层也是对的
        node.setHeight(1 + Math.max(getHeight(node.getLeftChild()), getHeight(node.getRightChild())));

        // 如果插入了新的节点，则需要进行平衡
        int balanceFactor = getBalanceFactor(node);
        if (balanceFactor > 1) {// L 类型
            AVLTreeNode<K> left = node.getLeftChild();
            // 左子节点的左子树高于右子树 LL类型，右旋
            if (getBalanceFactor(left) >= 0) {
                return rightRotate(node);
            } else {// 否则左子节点的右子树高于左子树，LR类型，先将左子节点的左子节点左旋，然后将节点右旋
                node.setLeftChild(leftRotate(left));
                return rightRotate(node);
            }
        } else if (balanceFactor < -1) {// R 类型
            AVLTreeNode<K> right = node.getRightChild();
            // 右子节点的右子树高度高于左子树，RR类型，节点左旋
            if (getBalanceFactor(right) <= 0) {
                return leftRotate(node);
            } else {// 右子节点的左子树高于右子树，RL类型，先将右子节点的右子节点右旋，然后将节点左旋
                node.setRightChild(rightRotate(right));
                return leftRotate(node);
            }
        } else {// 平衡，则不需要对子节点进行旋转
            return node;
        }
    }

    /**
     * 获取节点高度，规定节点不存在时，高度为0
     */
    private int getHeight(AVLTreeNode<K> node) {
        if (node == null) {
            return 0;
        } else {
            return node.getHeight();
        }
    }

    /**
     * 获取平衡因子，平衡因子 = 左子树高度 - 右子树高度
     */
    private int getBalanceFactor(AVLTreeNode<K> node) {
        if (node == null) {
            return 0;
        } else {
            return getHeight(node.getLeftChild()) - getHeight(node.getRightChild());
        }
    }

    /**
     *     A                    B
     *    / \                  / \
     *   B   C   右旋之后为    D   A
     *  / \                     / \
     * D   E                   E   C
     *
     * C D E 节点可能不存在，C D节点与父节点之间关系不变，在代码中不体现，为空不需要判断
     * D节点的父节点发生改变，在设置D节点的父节点指针时需要先判断D是否存在
     *
     * 这里，C D E 在旋转前后都是叶子节点，则说明C D E与其子节点关系不变，则高度不变
     * 而A B的子节点发生了改变，需要重新计算A B的高度，而且旋转后A是B的右子节点，需要先计算A的高度，再计算B的高度
     *
     */
    private AVLTreeNode<K> rightRotate(AVLTreeNode<K> node) {
        AVLTreeNode<K> A = node;
        if (A == null) {
            return null;
        }

        AVLTreeNode<K> B = A.getLeftChild();
        AVLTreeNode<K> E = B.getRightChild();

        // 将A挂到B下
        B.setParent(A.getParent());
        B.setRightChild(A);

        // 将E挂到A下
        A.setLeftChild(E);

        // 重新计算高度
        A.setHeight(Math.max(getHeight(A.getLeftChild()), getHeight(A.getRightChild())) + 1);
        B.setHeight(Math.max(getHeight(B.getLeftChild()), getHeight(B.getRightChild())) + 1);

        return B;
    }

    /**
     *     A                        C
     *    / \                      / \
     *   B   C    左旋之后为       A   E
     *      / \                 / \
     *     D  E                B   D
     *
     * B D E 节点有可能不存在，但是B E与其父节点之间关系不变，为空不需要判断
     * D节点的父节点发生改变，因此在D父节点指针的时候需要判断D是否存在
     *
     * 这里，B D E 在旋转前后都是叶子节点，则说明B D E与其子节点关系不变，则高度不变
     * 而A C的子节点发生了改变，需要重新计算A C的高度，而且旋转后A是C的左子节点，需要先计算A的高度，再计算C的高度
     *
     */
    private AVLTreeNode<K> leftRotate(AVLTreeNode<K> node) {
        AVLTreeNode<K> A = node;
        if (A == null) {
            return null;
        }

        // 需保证，左旋节点必须有右子节点
        AVLTreeNode<K> C = A.getRightChild();
        AVLTreeNode<K> D = C.getLeftChild();

        // 将A挂到C下面
        C.setParent(A.getParent());
        C.setLeftChild(A);

        // 将D挂到A下面
        A.setRightChild(D);

        // 重新计算高度
        A.setHeight(Math.max(getHeight(A.getLeftChild()), getHeight(A.getRightChild())) + 1);
        C.setHeight(Math.max(getHeight(C.getLeftChild()), getHeight(C.getRightChild())) + 1);

        return C;
    }

    @Override
    public K search(K key) {
        return searchNode(this.root, key).getKey();
    }

    private AVLTreeNode<K> searchNode(AVLTreeNode<K> node, K key) {
        int compareResult;
        if ((compareResult = node.getKey().compareTo(key)) < 0) {
            return searchNode(node.getRightChild(), key);
        } else if (compareResult > 0) {
            return searchNode(node.getLeftChild(), key);
        } else {
            return node;
        }
    }

    public void inOrderWalk(Consumer<K> consumer) {
        RecursionBinaryTreeWalker.inOrderTraversal(this.root, consumer);
    }

}
